'''
快速增强点中，每隔一定小时数是一个颜色
'''

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,split_str_id,tydat_CMA
from geopy.distance import geodesic
import cartopy.io.shapereader as shpreader
import geopandas as gpd



# names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
names = ['dusurui_16']
tynames,tyids = split_str_id(names)
drawopt = True


for ty,tyid in zip(tynames,tyids):

    RIstd = 7
    dt = timedelta(hours=48)  # 🔧 改为12小时间隔，可以根据需要调整
    
    directory =  f'D:\\met_data\\ty_ensemble\\{ty}_{tyid}\\'    # 要筛选的文件开头
    dates_name = os.listdir(directory)
    dates_paths = [directory + date for date in dates_name]
    obs_path = rf'D:\met_data\ty_obs\{ty}_CMAobs.txt'
    pic_savepath = rf'C:\Users\lenovo\Desktop\typic\RIpoints\{ty}'
    os.makedirs(pic_savepath,exist_ok=True)
    
    def basemap(ax, title=None, extend_degrees=15):
        """
        只显示中国省界的底图
        
        推荐使用: admin_1_states_provinces (面状数据)
        """
        
        ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.5)
        
        # ✅ 使用 admin_1_states_provinces (面状数据，更完整)
        try:
            provinces_shp = shpreader.natural_earth(
                resolution='10m',
                category='cultural',
                name='admin_1_states_provinces'  # 面状，不是 _lines
            )
            
            gdf = gpd.read_file(provinces_shp)
            
            # 筛选中国省份
            china_provinces = gdf[
                (gdf['admin'] == 'China') | 
                (gdf['adm0_a3'] == 'CHN')
            ]
            
            # 绘制省界
            ax.add_geometries(
                china_provinces.geometry,
                ccrs.PlateCarree(),
                facecolor='none',
                edgecolor='black',
                linewidth=0.5
            )
            
            print(f"✅ 成功绘制 {len(china_provinces)} 个中国省份")
            
        except Exception as e:
            print(f"❌ 加载失败: {e}")
        
        # 网格线
        gl = ax.gridlines(
            draw_labels=True, 
            alpha=0.3, 
            linestyle='--', 
            linewidth=0.5,
            color='gray'
        )
        
        gl.top_labels = False
        gl.right_labels = False
        gl.xlabel_style = {'size': 8}
        gl.ylabel_style = {'size': 8}
        gl.xformatter = LongitudeFormatter(number_format='.1f')
        gl.yformatter = LatitudeFormatter(number_format='.1f')
        
        if title:
            ax.set_title(title)
    
    def track(ty, ax, lon, lat, title=None, color=None, marker=".", label=None, show_label=False, s=12, time_data=None):
        """
        绘制台风轨迹（时间分段显示）
        
        Parameters:
        time_data: 如果提供，使用此时间数据；否则使用 ty.time
        """

        
        # 确定使用的时间数据
        if time_data is not None:
            time_array = np.array(time_data)
        else:
            time_array = ty.time
        
        # 🔧 按时间间隔分段绘制
        now = start
        i = 0
        colors = plt.cm.tab20(np.linspace(0, 1, 20))  # 增加颜色数量
        
        lon = np.array(lon)
        lat = np.array(lat)
        
        # 🔧 计算总时间跨度，确定循环次数
        if len(time_array) > 0:
            total_duration = max(time_array) - start
            max_iterations = int(total_duration / dt) + 2  # 防止无限循环
        else:
            max_iterations = 20  # 默认最大迭代次数
        
        iteration = 0
        while now <= max(time_array) and iteration < max_iterations:
            end = now + dt
            
            # 🔧 根据时间筛选点
            if time_data is not None:
                want = (time_array >= now) & (time_array < end)
                if hasattr(ty, 'time') and len(ty.time) > 0:
                    # 如果还有record信息，也需要考虑
                    try:
                        record_mask = (record == 1)
                        want = want & record_mask
                    except:
                        pass  # 如果没有record，只按时间筛选
            else:
                want = (time_array >= now) & (time_array < end)
                try:
                    want = want & (record == 1)
                except:
                    pass
            
            # 🔧 生成时间标签（按小时显示）
            if show_label:
                hours_from_start = int((now - start).total_seconds() / 3600)
                hours_to_end = int((end - start).total_seconds() / 3600)
                if dt.total_seconds() >= 24 * 3600:  # 如果间隔>=1天，显示天数
                    label_to_use = f"{dt.days*i}-{dt.days*(i+1)}d"
                else:  # 否则显示小时
                    label_to_use = f"{hours_from_start}h-{hours_to_end}h"
            else:
                label_to_use = None
            
            # 绘制当前时间段的点
            if np.any(want):
                ax.scatter(
                    lon[want], lat[want], 
                    marker=marker, 
                    transform=ccrs.PlateCarree(), 
                    color=colors[i % len(colors)],  # 防止颜色索引越界
                    label=label_to_use, 
                    s=s,
                    alpha=0.8,
                    edgecolors='white',
                    linewidths=0.5
                )
            
            now = end
            i += 1
            iteration += 1
    
    # 🔧 修改后的 track_of_members 函数
    def track_of_members(ax, ty, time, lat, lon, marker='.', show_label=True, s=24):
        """
        绘制集合成员轨迹
        """

        
        # 🔧 按时间间隔分段绘制
        now = start
        i = 0
        colors = plt.cm.tab20(np.linspace(0, 1, 20))
        time = np.array(time)
        lat = np.array(lat)
        lon = np.array(lon)
        
        # 计算最大时间范围
        if len(time) > 0:
            max_time = min(start + timedelta(days=16), max(time))
        else:
            max_time = start + timedelta(days=16)
        
        iteration = 0
        max_iterations = int((max_time - start) / dt) + 2
        
        while now <= max_time and iteration < max_iterations:
            end = now + dt
            want = (time >= now) & (time < end)
            
            # 🔧 生成时间标签
            if show_label:
                hours_from_start = int((now - start).total_seconds() / 3600)
                hours_to_end = int((end - start).total_seconds() / 3600)
                if dt.total_seconds() >= 24 * 3600:  # 如果间隔>=1天
                    label_to_use = f"{dt.days*i}-{dt.days*(i+1)}d"
                else:  # 小时间隔
                    label_to_use = f"{hours_from_start}h-{hours_to_end}h"
            else:
                label_to_use = None
            
            # 绘制当前时间段的点
            if np.any(want):
                ax.scatter(
                    lon[want], lat[want], 
                    marker=marker, 
                    transform=ccrs.PlateCarree(), 
                    color=colors[i % len(colors)], 
                    label=label_to_use, 
                    s=s,
                    alpha=0.7,
                    edgecolors='white',
                    linewidths=0.5
                )
            
            now = end
            i += 1
            iteration += 1
    
    # 起报时间循环
    for i in tqdm.tqdm(range(len(dates_paths)), desc=f"Processing {ty}"):
        members_paths = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f != "TRACK_ID_0" and f.startswith("TRACK_ID")]
        era_path = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f.startswith("TRACK_ID_0")][0]
        
        fig, ax = plt.subplots(1, 1, figsize=(12, 10), subplot_kw={'projection': ccrs.PlateCarree()})
        start = datetime.strptime(dates_name[i], "%Y%m%d%H")
        
        print(f"处理起报时间: {start}")
        
        '''draw members'''
        # 循环 某一起报时间的成员
        time = []
        lat = []
        lon = []
        
        for j in range(len(members_paths)):
            try:
                member = tydat(members_paths[j], RIstd, lon_format='180')
                record = member.num_rapidgrow()
                RI = (record == 1)
                if RI.any():
                    time.extend(member.time[RI].tolist())
                    lat.extend(member.lat[RI].tolist())
                    lon.extend(member.lon[RI].tolist())
            except Exception as e:
                print(f"处理成员 {j} 时出错: {e}")
                continue
        
        print(f"集合成员RI点数量: {len(time)}")
        
        # 绘制集合成员RI点
        if len(time) > 0:
            track_of_members(ax, member, time, lat, lon, marker="o", show_label=True, s=64)
        
        ''' draw Analysis '''
        try:
            era = tydat(era_path, RIstd, lon_format='180')
            record = era.num_rapidgrow()
            track(era, ax, era.lon, era.lat, marker="s", s=96, show_label=False, time_data=era.time)
        except Exception as e:
            print(f"处理分析数据时出错: {e}")
        
        ''' obs '''
        try:
            obs = tydat_CMA(obs_path, lon_format='180')
            record = count_rapidgrow(RIstd, obs.umax, obs.time)
            track(obs, ax, obs.lon, obs.lat, marker="^", s=96, show_label=False, time_data=obs.time)
        except Exception as e:
            print(f"处理观测数据时出错: {e}")
        
        basemap(ax)
        
        
        # 设置标题和图例
        title = dates_name[i] + f" {ty} - RI Points by Time"
        ax.set_title(title, fontsize=14, pad=20)
        
        if ty == 'dusurui' :
            print(ty)
            ax.set_extent([115,130,15,25] , crs=ccrs.PlateCarree())        
            
        # 🔧 改进图例显示
        ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=10)
        # 添加说明文字
        ax.text(x=0.02, y=0.98, s="○ Ensemble RI", transform=ax.transAxes, 
                color="black", fontsize=10, va='top',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        ax.text(x=0.02, y=0.90, s="□ Analysis RI", transform=ax.transAxes, 
                color="black", fontsize=10, va='top',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        ax.text(x=0.02, y=0.80, s="△ Observation RI", transform=ax.transAxes, 
                color="black", fontsize=10, va='top',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        
        
        plt.tight_layout()
        if drawopt == False: 
            plt.savefig(f'{pic_savepath}\\{title.replace(" ", "_")}.png', 
                       dpi=300, bbox_inches='tight')
            plt.close()
        else:
            plt.show()
        
        print(f"完成 {title}")